﻿using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using Tensorflow;
using Tensorflow.NumPy;
using static Tensorflow.Binding;

namespace TensorFlowTest.BasicModels;

public class KMeansClustering : SciSharpExample, IExample
{
    public int? train_size = null;
    public int validation_size = 5000;
    public int? text_size = null;
    public int batch_size = 1024;//the number of samples per batch

    Datasets<MnistDataSet> mnist;

    NDArray full_data_x;
    int num_steps = 20;//total steps to train
    int k = 25;//the number of clusters
    int num_classes = 10;//the 10 digits


    float accuray_test = 0f;


    public ExampleConfig InitConfig()
    {
        Config= new ExampleConfig()
        {
            Name = "K-means Clustering",
            Enabled = false,
            IsImportingGraph = true,
        };
        return Config;
    }

    public bool Run()
    {
        PrepareData();
        var graph = ImportGraph();
        using (var sess = tf.Session(graph))
        {
            Train(sess);
        }
        return accuray_test > 0.70;
    }

    public override void PrepareData()
    {
        var loader = new MnistModelLoader();
        var setting = new ModelLoadSetting
        {
            TrainDir = ".resources/mnist",
            OneHot = true,
            TrainSize = train_size,
            ValidationSize = validation_size,
            TestSize = text_size,
            ShowProgressInConsole = true
        };
        mnist = loader.LoadAsync(setting).Result;
        full_data_x = mnist.Train.Data;

        // download graph meta data
        string url = "https://raw.githubusercontent.com/SciSharp/TensorFlow.NET/master/graph/kmeans.meta";
        loader.DownloadAsync(url, ".resources/graph", "kmeans.meta").Wait();
    }

    public override Graph ImportGraph()
    {
        var graph = tf.Graph().as_default();
        tf.train.import_meta_graph(".resources/graph/kmeans.meta");
        return graph;
    }

    private void Train(Session sess)
    {
        var graph = sess.graph;

        //input images
        Tensor X = graph.get_operation_by_name("Placeholder");
        // Labels
        Tensor Y = graph.get_operation_by_name("Placeholder_1");

        //k-means param
        var init_vars = tf.global_variables_initializer();
        Tensor init_op = graph.get_operation_by_name("cond/Merge");
        var train_op = graph.get_operation_by_name("group_deps");
        Tensor avg_distance = graph.get_operation_by_name("Mean");
        Tensor cluster_idx = graph.get_operation_by_name("Squeeze_1");
        NDArray[] result = null;

        sess.run(init_vars, new FeedItem(X, full_data_x));
        sess.run(init_op, new FeedItem(X, full_data_x));

        //training 
        var sw = new Stopwatch();

        foreach (var i in range(1, num_steps + 1))
        {
            sw.Restart();
            result = sess.run(
                new ITensorOrOperation[]
                {
                train_op,avg_distance,cluster_idx
                },
                new FeedItem(X, full_data_x)
            );
            sw.Stop();

            if (i % 4 == 0 || i == 1)
            {
                print($"Step {i}, Avg Distance: {result[1]} Elapse:{sw.ElapsedMilliseconds}ms");
            }

        }


        var idx = result[2].ToArray<Int64>();

        var counts = np.zeros((k, num_classes), np.float32);

        sw.Start();
        foreach (var i in range(idx.Length))
        {
            var x = mnist.Train.Labels[i];
            //counts[idx[i]] += x;
            throw new NotImplementedException("");
        }
        sw.Stop();
        print($"Assign a label to each centroid took {sw.ElapsedMilliseconds}ms");

        var labels_map_array = np.argmax(counts, 1);
        var labels_map = tf.convert_to_tensor(labels_map_array);


        var cluster_label = tf.nn.embedding_lookup(labels_map, cluster_idx);

        //compute accuracy
        var correct_prediction = tf.equal(cluster_label, tf.cast(tf.math.argmax(Y, 1), tf.int32));

        var cast = tf.cast(correct_prediction, tf.float32);
        var accuracy_op = tf.reduce_mean(cast);


        //test model
        var (test_x, text_y) = (mnist.Test.Data, mnist.Test.Labels);
        print($"Test Accuracy:{accuray_test}");

    }
}
